from langchain.chains import RetrievalQA
from langchain_community.vectorstores import FAISS
from langchain_openai import ChatOpenAI
from openai import OpenAI
import os

class RAGChain:
    def __init__(self, api_key: str, api_base: str, vector_store: FAISS, top_k: int):
        # 创建 OpenAI 客户端
        client = OpenAI(
            api_key=api_key,
            base_url=api_base
        )
        
        # 创建 LangChain LLM
        self.llm = ChatOpenAI(
            openai_api_key=api_key,
            openai_api_base=api_base,
            model="deepseek-chat",
            temperature=0.7,
            max_tokens=2000
        )
        
        # 创建检索器
        retriever = vector_store.as_retriever(search_kwargs={"k": top_k})
        
        # 创建 RAG 链
        self.qa_chain = RetrievalQA.from_chain_type(
            llm=self.llm,
            chain_type="stuff",
            retriever=retriever,
            return_source_documents=True
        )
    
    def query(self, question: str) -> dict:
        """查询知识库"""
        try:
            return self.qa_chain.invoke({"query": question})
        except Exception as e:
            print(f"查询出错: {str(e)}")
            return {"result": "查询失败", "source_documents": []} 